
#ifndef VESSEL_MASK_H
#define VESSEL_MASK_H

#include "dsaio.h"

#include "itkBinaryThresholdImageFilter.h"
#include "itkCastImageFilter.h"
#include "itkHessian2DToVesselnessMeasureImageFilter.h"
#include "itkHessianRecursiveGaussianImageFilter.h"
#include "itkImageRegionConstIterator.h"
#include "itkImageRegionIterator.h"
#include "itkInvertIntensityImageFilter.h"
#include "itkRescaleIntensityImageFilter.h"
#include "itkSymmetricSecondRankTensor.h"

#define UCHAR_THRESHOLD 10

/**
 * Returns vessel mask of the input image. 0 = background, other = vessel.
 */
DsaFrame::Pointer vessel_mask ( DsaFrame::Pointer input )
{
  vector< double > radius_pyramid;
  radius_pyramid.push_back ( 1.0 );
  radius_pyramid.push_back ( 1.5 );
  radius_pyramid.push_back ( 2.0 );
  radius_pyramid.push_back ( 3.0 );
  radius_pyramid.push_back ( 4.0 );
  radius_pyramid.push_back ( 5.0 );
  
  // invert to make vessels light, bg dark
  typedef itk::InvertIntensityImageFilter< DsaFrame, DsaFrame > Invert;
  Invert::Pointer invert = Invert::New();
  invert->SetInput( input );
  invert->Update();
    
  // cast pixels to double for hessian filter
  typedef double DoublePixel;
  typedef itk::Image< DoublePixel, 2 > DoubleImage;
  typedef itk::CastImageFilter< DsaFrame, DoubleImage > CastToDouble;
  CastToDouble::Pointer castToDouble = CastToDouble::New();
  castToDouble->SetInput( invert->GetOutput() );
  castToDouble->Update();
  // cerr << "casted to double" << endl;
  
  // vesselness filter setup
  typedef float FloatPixel;
  typedef itk::Image< FloatPixel, 2 > FloatImage;
  typedef itk::HessianRecursiveGaussianImageFilter< DoubleImage >
          HessianFilter;
  typedef itk::Hessian2DToVesselnessMeasureImageFilter< FloatPixel >
          VesselnessFilter;
  HessianFilter::Pointer hessianFilter = HessianFilter::New();
  VesselnessFilter::Pointer vesselnessFilter = VesselnessFilter::New();
  hessianFilter->SetInput ( castToDouble->GetOutput() );
  hessianFilter->SetSigma ( radius_pyramid[ 0 ] );
  hessianFilter->Update();
  vesselnessFilter->SetInput ( hessianFilter->GetOutput() );
  vesselnessFilter->Update();
  
  // create output image
  FloatImage::Pointer maxImage = FloatImage::New();
  maxImage->SetRegions ( input->GetLargestPossibleRegion() );
  maxImage->Allocate();
  
  // define iterators
  typedef itk::ImageRegionIterator< FloatImage > FloatImageIterator;
  FloatImageIterator maxIt ( maxImage, maxImage->GetLargestPossibleRegion() );
  typedef itk::ImageRegionConstIterator< FloatImage > FloatImageConstIterator;
  FloatImageConstIterator vnessIt ( vesselnessFilter->GetOutput(),
    vesselnessFilter->GetOutput()->GetLargestPossibleRegion() );
  
  // fill max image with initial values
  for ( maxIt.GoToBegin(), vnessIt.GoToBegin(); ! maxIt.IsAtEnd();
        ++ maxIt, ++ vnessIt )
  {
    maxIt.Set ( vnessIt.Get() );
  }
  
  for ( int i = 1; i < radius_pyramid.size(); ++ i )
  {
    double sigma = radius_pyramid[ i ];
    cerr << "  sigma = " << sigma;
    
    // vesselnessFilter->ResetPipeline();
    hessianFilter->SetSigma ( sigma );
	vesselnessFilter->Update();
    
    int counter = 0;
    int total = 0;    
    for ( vnessIt.GoToBegin(), maxIt.GoToBegin(); ! vnessIt.IsAtEnd();
          ++ vnessIt, ++ maxIt )
    {
      float current = vnessIt.Get();
      float max = maxIt.Get();
      if ( max < current )
      {
        maxIt.Set( current );
        ++ counter;
      }
      ++ total;
    }
    cerr << " - updated " << counter << " pixels of "
         << total << " pixels" << endl;
  }
  
  // scale back between 0 and 255
  typedef itk::RescaleIntensityImageFilter< FloatImage, FloatImage>
    RescaleFilter;
  RescaleFilter::Pointer rescaleFilter = RescaleFilter::New();
  rescaleFilter->SetInput ( maxImage );
  rescaleFilter->SetOutputMinimum ( 0 );
  rescaleFilter->SetOutputMaximum ( 255 );
  rescaleFilter->Update();
  
  // cast pixels back to unsigned char
  typedef itk::CastImageFilter< FloatImage, DsaFrame > CastToUChar;
  CastToUChar::Pointer castToUChar = CastToUChar::New();
  castToUChar->SetInput ( rescaleFilter->GetOutput() );
  castToUChar->Update();
  
  /* This part commented so threshold level can be set later
  // thresholding
  typedef itk::BinaryThresholdImageFilter< DsaFrame, DsaFrame > Threshold;
  Threshold::Pointer threshold = Threshold::New();
  threshold->SetInput ( castToUChar->GetOutput() );
  threshold->SetInsideValue ( 255 );
  threshold->SetOutsideValue ( 0 );
  threshold->SetLowerThreshold ( UCHAR_THRESHOLD ); // magic number
  threshold->SetUpperThreshold ( 255 );
  threshold->Update();
  */
  // return threshold->GetOutput();
  
  return castToUChar->GetOutput();
}

#endif
